package com.ceair.authorization.federation;

import com.ceair.constant.SecurityConstants;
import com.ceair.entity.Oauth2BasicUser;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.oauth2.core.oidc.IdTokenClaimNames;
import org.springframework.security.oauth2.core.oidc.OidcIdToken;
import org.springframework.security.oauth2.core.oidc.endpoint.OidcParameterNames;
import org.springframework.security.oauth2.core.oidc.user.OidcUser;
import org.springframework.security.oauth2.core.user.OAuth2User;
import org.springframework.security.oauth2.jwt.JwtClaimsSet;
import org.springframework.security.oauth2.server.authorization.token.JwtEncodingContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenCustomizer;

import java.util.*;
import java.util.stream.Collectors;

/**
 * @author wangbaohai
 * @ClassName FederatedIdentityIdTokenCustomizer
 * @description: 联合身份认证自定义token处理，当使用openId Connect登录时将用户信息写入idToken中
 * @date 2024年11月26日
 * @version: 1.0.0
 */
public class FederatedIdentityIdTokenCustomizer implements OAuth2TokenCustomizer<JwtEncodingContext> {

    /**
     * 存储ID令牌中包含的全部声明名称
     * 这些声明是OpenID Connect标准的一部分，用于描述用户身份和认证相关信息
     * 使用不可变集合Set来确保这些声明名称不会被修改，从而增强安全性
     */
    private static final Set<String> ID_TOKEN_CLAIMS = Set.of(
            IdTokenClaimNames.ISS,  // 发行者标识
            IdTokenClaimNames.SUB,  // 主体标识（用户唯一标识）
            IdTokenClaimNames.AUD,  // 接收者标识
            IdTokenClaimNames.EXP,  // 过期时间
            IdTokenClaimNames.IAT,  // 发行时间
            IdTokenClaimNames.AUTH_TIME,  // 认证时间
            IdTokenClaimNames.NONCE,  // 用于防止重放攻击的一次性值
            IdTokenClaimNames.ACR,  // 认证上下文类参考
            IdTokenClaimNames.AMR,  // 认证方法引用
            IdTokenClaimNames.AZP,  // 授权方标识
            IdTokenClaimNames.AT_HASH,  // 访问令牌散列值
            IdTokenClaimNames.C_HASH  // 代码散列值
    );

    /**
     * 根据JwtEncodingContext上下文定制JWT
     * 该方法主要用于在JWT中添加自定义声明
     *
     * @param context JwtEncodingContext实例，包含编码JWT所需的信息
     */
    @Override
    public void customize(JwtEncodingContext context) {
        // 当令牌类型为ID_TOKEN时，执行自定义逻辑
        if (OidcParameterNames.ID_TOKEN.equals(context.getTokenType().getValue())) {
            // 从用户主体中提取第三方声明
            Map<String, Object> thirdPartyClaims = extractClaims(context.getPrincipal());
            // 移除已存在的声明，避免重复，并添加第三方声明到JWT中
            context.getClaims().claims(existingClaims -> {
                existingClaims.keySet().forEach(thirdPartyClaims::remove);

                ID_TOKEN_CLAIMS.forEach(thirdPartyClaims::remove);

                existingClaims.putAll(thirdPartyClaims);
            });
        }

        // 检查登录用户信息是不是OAuth2User，在token中添加loginType属性
        if (context.getPrincipal().getPrincipal() instanceof OAuth2User user) {
            JwtClaimsSet.Builder claims = context.getClaims();
            Object loginType = user.getAttribute("loginType");
            if (loginType instanceof String) {
                // 同时检验是否为String和是否不为空
                claims.claim("loginType", loginType);
            }

            // 获取申请的scopes
            Set<String> scopes = context.getAuthorizedScopes();
            // 获取用户的权限
            Collection<? extends GrantedAuthority> authorities = user.getAuthorities();
            // 提取权限并转为字符串
            Set<String> authoritySet = Optional.ofNullable(authorities).orElse(Collections.emptyList()).stream()
                    // 获取权限字符串
                    .map(GrantedAuthority::getAuthority)
                    // 去重
                    .collect(Collectors.toSet());
            // 合并scope与用户信息
            authoritySet.addAll(scopes);

            // 将权限信息放入jwt的claims中（也可以生成一个以指定字符分割的字符串放入）
            claims.claim(SecurityConstants.AUTHORITIES_KEY, authoritySet);
            claims.claim(SecurityConstants.TOKEN_UNIQUE_ID, Objects.requireNonNull(user.getAttribute("uniqueId")));
        }

        // 检查登录用户信息是不是UserDetails，排除掉没有用户参与的流程
        if (context.getPrincipal().getPrincipal() instanceof Oauth2BasicUser user) {
            // 获取申请的scopes
            Set<String> scopes = context.getAuthorizedScopes();
            // 获取用户的权限
            Collection<? extends GrantedAuthority> authorities = user.getAuthorities();
            // 提取权限并转为字符串
            Set<String> authoritySet = Optional.ofNullable(authorities)
                    .orElse(Collections.emptyList()).stream()
                    // 获取权限字符串
                    .map(GrantedAuthority::getAuthority)
                    // 去重
                    .collect(Collectors.toSet());

            // 合并scope与用户信息
            authoritySet.addAll(scopes);

            JwtClaimsSet.Builder claims = context.getClaims();
            // 将权限信息放入jwt的claims中（也可以生成一个以指定字符分割的字符串放入）
            claims.claim(SecurityConstants.AUTHORITIES_KEY, authoritySet);
            claims.claim(SecurityConstants.TOKEN_UNIQUE_ID, user.getId());
            // 放入其它自定内容
            // 角色、头像...
            claims.claim("avatarUrl", user.getAvatarUrl());
        }
    }

    /**
     * 从Authentication对象中提取用户声明（claims）
     * 该方法处理的Authentication对象可以包含不同类型的用户主体，如OidcUser或OAuth2User
     * 根据用户主体的类型，提取相应的声明信息
     *
     * @param principal Authentication对象，包含用户的身份信息
     * @return 包含用户声明的Map对象，如果无法提取声明，则返回空Map
     */
    private Map<String, Object> extractClaims(Authentication principal) {
        Map<String, Object> claims;

        // 判断用户主体是否为OidcUser类型，如果是，则从其ID令牌中提取声明
        if (principal.getPrincipal() instanceof OidcUser oidcUser) {
            OidcIdToken idToken = oidcUser.getIdToken();
            claims = idToken.getClaims();
        } else if (principal.getPrincipal() instanceof OAuth2User oauth2User) {
            // 如果用户主体是OAuth2User类型，则直接提取其属性作为声明
            claims = oauth2User.getAttributes();
        } else {
            // 如果用户主体类型不匹配，则返回空Map
            claims = Collections.emptyMap();
        }

        // 返回一个新的HashMap对象，包含提取到的声明
        return new HashMap<>(claims);
    }

}
